import gzip

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 检查是否有可用的 GPU，如果没有则使用 CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")


# 定义生成器和判别器
class Generator(nn.Module):
    def __init__(self, z_dim=100, img_shape=(1, 28, 28)):
        super(Generator, self).__init__()
        self.img_shape = img_shape
        self.gen = nn.Sequential(
            nn.Linear(z_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def forward(self, x):
        x = self.gen(x)
        x = x.view(-1, *self.img_shape)
        return x


class Discriminator(nn.Module):
    def __init__(self, img_shape=(1, 28, 28)):
        super(Discriminator, self).__init__()
        self.img_shape = img_shape
        self.disc = nn.Sequential(
            nn.Linear(int(np.prod(img_shape)), 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = x.view(-1, int(np.prod(self.img_shape)))
        x = self.disc(x)
        return x


# 定义加载本地 MNIST 数据集的函数
MNIST_FILE_PATH = 'D:/TT_WORK+/PyCharm/20250109_1_CNN/MNIST/'


def load_data():
    # 加载图像数据
    with gzip.open(MNIST_FILE_PATH + 'train-images-idx3-ubyte.gz', 'rb') as f:  # 训练集
        X_train = np.frombuffer(f.read(), dtype=np.uint8, offset=16).reshape(-1, 28 * 28)

    with gzip.open(MNIST_FILE_PATH + 't10k-images-idx3-ubyte.gz', 'rb') as f:  # 测试集
        X_test = np.frombuffer(f.read(), dtype=np.uint8, offset=16).reshape(-1, 28 * 28)

    # 加载标签数据
    with gzip.open(MNIST_FILE_PATH + 'train-labels-idx1-ubyte.gz', 'rb') as f:  # 训练集标签
        y_train = np.frombuffer(f.read(), dtype=np.uint8, offset=8)

    with gzip.open(MNIST_FILE_PATH + 't10k-labels-idx1-ubyte.gz', 'rb') as f:  # 测试集标签
        y_test = np.frombuffer(f.read(), dtype=np.uint8, offset=8)

    return (X_train, y_train), (X_test, y_test)


# 加载数据并转换为 PyTorch 张量
(X_train, y_train), (X_test, y_test) = load_data()

# 将数据转换为 PyTorch 张量并归一化到 [-1, 1] 范围
X_train = torch.tensor(X_train, dtype=torch.float32).view(-1, 1, 28, 28) / 255.0 * 2 - 1
X_test = torch.tensor(X_test, dtype=torch.float32).view(-1, 1, 28, 28) / 255.0 * 2 - 1

# 创建数据集和数据加载器
train_dataset = TensorDataset(X_train)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 初始化生成器和判别器，并将它们移动到设备上
z_dim = 100
img_dim = (1, 28, 28)  # 确保图像的维度是 (channels, height, width)

generator = Generator(z_dim, img_dim).to(device)
discriminator = Discriminator(img_dim).to(device)

lr = 0.0001
num_epochs = 100

optimizer_gen = optim.Adam(generator.parameters(), lr=lr)
optimizer_dis = optim.Adam(discriminator.parameters(), lr=lr)

criterion = nn.BCELoss()

# 记录损失值
train_loss_g = []
train_loss_d = []

# 训练过程
for epoch in range(num_epochs):
    gen_loss_epoch = 0
    disc_loss_epoch = 0

    for batch_idx, (real,) in enumerate(train_loader):
        real = real.to(device)
        batch_size = real.size(0)

        # 训练判别器
        noise = torch.randn(batch_size, z_dim, device=device)
        fake = generator(noise)
        disc_real_loss = criterion(discriminator(real), torch.ones(batch_size, 1, device=device))
        disc_fake_loss = criterion(discriminator(fake.detach()), torch.zeros(batch_size, 1, device=device))
        disc_loss = (disc_real_loss + disc_fake_loss) / 2

        optimizer_dis.zero_grad()
        disc_loss.backward()
        optimizer_dis.step()

        # 训练生成器
        noise = torch.randn(batch_size, z_dim, device=device)
        fake = generator(noise)
        gen_loss = criterion(discriminator(fake), torch.ones(batch_size, 1, device=device))

        optimizer_gen.zero_grad()
        gen_loss.backward()
        optimizer_gen.step()

        gen_loss_epoch += gen_loss.item()
        disc_loss_epoch += disc_loss.item()

        if batch_idx % 100 == 0:
            print(f"Epoch [{epoch}/{num_epochs}] Batch {batch_idx}/{len(train_loader)} \
                  Loss D: {disc_loss.item():.4f}, loss G: {gen_loss.item():.4f}")

    # 记录每个 epoch 的平均损失
    train_loss_g.append(gen_loss_epoch / len(train_loader))
    train_loss_d.append(disc_loss_epoch / len(train_loader))

    # 每个 epoch 保存一些生成的图像
    generator.eval()
    with torch.no_grad():
        noise = torch.randn(1, z_dim, device=device)
        generated_img = generator(noise).view(28, 28).cpu().numpy()
        plt.imshow(generated_img, cmap='gray')
        plt.savefig(f'generated_img_epoch_{epoch}.png')
        plt.close()
    generator.train()

# 5-保存模型
torch.save(generator.state_dict(), 'generator.pt')
torch.save(discriminator.state_dict(), 'discriminator.pt')

# 绘制训练损失曲线
plt.figure(figsize=(10, 5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(train_loss_g, label="G")
plt.plot(train_loss_d, label="D")
plt.xlabel("Epochs")
plt.ylabel("Loss")
plt.legend()
plt.savefig("GAN_loss_curve.png")
plt.show()
